//originally from http://www.javablogging.com/dynamic-in-memory-compilation/

package compiler;

import java.io.IOException;
import java.security.SecureClassLoader;
import java.util.ArrayList;
import java.util.Collection;
import java.util.HashMap;
import java.util.List;
import javax.tools.FileObject;
import javax.tools.ForwardingJavaFileManager;
import javax.tools.JavaFileObject;
import javax.tools.StandardJavaFileManager;

public class ClassFileManager extends
		ForwardingJavaFileManager {
    /**
     * Maps the class names to its compiled code object.
     */
    private HashMap<String, JavaClassObject> classNamesMap = new HashMap<String, JavaClassObject>();
    
    /**
     * For addition of classes after loading them
     */
    private ClassLoader classLoader = new SecureClassLoader() {
        @Override
        protected Class<?> findClass(String name)
            throws ClassNotFoundException {
            JavaClassObject jclassObject = classNamesMap.get(name);
            
            if (jclassObject == null)
                throw new ClassNotFoundException("class " + name + " is not found.");

            byte[] b = jclassObject.getBytes();
            return super.defineClass(name, jclassObject
                .getBytes(), 0, b.length);
        }
    };

    /**
    * Will initialize the manager with the specified
    * standard java file manager
    *
    * @param standardManger
    */
    public ClassFileManager(StandardJavaFileManager standardManager) 
    {
        super(standardManager);
    }
    
    public ClassLoader getClassLoader()
    {
        return classLoader;
    }
    
    public Collection<JavaClassObject> getClassObjects()
    {
        return classNamesMap.values();
    }
    
    /**
     * Obtains the compiled classes.
     * @return
     * @throws ClassNotFoundException 
     */    
    public List<Class<?>> getCompiledClasses() throws ClassNotFoundException
    {        
        List<Class<?>> result = new ArrayList<Class<?>>();
        
        for (String name : classNamesMap.keySet())
            result.add(classLoader.loadClass(name));
        
        return result;
    }

    /**
    * Gives the compiler an instance of the JavaClassObject
    * so that the compiler can write the byte code into it.
    */
    @Override
    public JavaFileObject getJavaFileForOutput(Location location,
        String className, JavaFileObject.Kind kind, FileObject sibling)
            throws IOException {
            JavaClassObject jclassObject = new JavaClassObject(className, kind);
            classNamesMap.put(className, jclassObject);
        return jclassObject;
    }
}
